# encoding=utf-8

import argparse
import datetime
import os
import time
import traceback

import cv2
import numpy as np

from NerveformerWrapper import NerveformerWrapper

nerveformer_infer = NerveformerWrapper()


def run_one(file_path: str):
    try:
        t1 = time.perf_counter()
        out = nerveformer_infer.infer(file_path)
        t2 = time.perf_counter()
        print("Infer file: {}, cost: {}s".format(os.path.basename(file_path), (t2 - t1)))
        return out
    except Exception as e:
        print(traceback.format_exc())
        return None


def generate_raw_data(from_path: str, raw_path: str):
    os.makedirs(raw_path, exist_ok=True)
    file_names = os.listdir(from_path)
    for each_name in file_names:
        abs_path = os.path.join(from_path, each_name)
        save_path = os.path.join(raw_path, each_name)

        out = run_one(abs_path)
        if out is not None:
            cv2.imwrite(save_path, out)


def generate_color_data(from_path: str, raw_path: str, color_path: str):
    os.makedirs(color_path, exist_ok=True)
    file_names = os.listdir(from_path)
    for each_name in file_names:
        t1 = time.perf_counter()
        abs_path = os.path.join(from_path, each_name)
        out_raw_path = os.path.join(raw_path, each_name)
        out_color_path = os.path.join(color_path, each_name)

        ori_file = cv2.imread(abs_path)
        raw_file = cv2.imread(out_raw_path)
        no_zero_index = np.nonzero(raw_file)
        ori_file[no_zero_index[0], no_zero_index[1], :] = \
            ori_file[no_zero_index[0], no_zero_index[1], :] * 0.5 + 0.5 * np.array([0, 212, 255])

        cv2.imwrite(out_color_path, ori_file)
        t2 = time.perf_counter()
        print("Draw file: {}, cost: {}s".format(each_name, (t2 - t1)))


def run(work_dir: str):
    from_path = os.path.join(work_dir, "user_input")
    raw_path = os.path.join(work_dir, "out_raw_data")
    merge_path = os.path.join(work_dir, "out_merge_data")

    if not os.path.exists(from_path):
        print("user_input dir is not found in {}, program will exit".format(work_dir))
        return

    print("Start to generate model raw data in: {}".format(raw_path))
    generate_raw_data(from_path, raw_path)

    print("Start to generate color data in: {}".format(merge_path))
    generate_color_data(from_path, raw_path, merge_path)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--workdir", type=str)
    args = parser.parse_args()
    if args.workdir == "" or args.workdir is None:
        print("work dir is empty")
        exit(-1)

    run(args.workdir)
